# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import functools
import logging
import os
import signal

from absl import app, flags
from tornado.ioloop import IOLoop, PeriodicCallback

import coin.flow.subscriber as flow_subscriber
from coin.base.param_util import to_list
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.exchange.shared.kr_rest.product import SharedProduct
from coin.exchange.util.feed_checker import get_default_feed_checker_config
from coin.flow.simple_queue import SimpleQueue
from coin.proto.coin_order_checker_pb2 import AccountCheckerConfig
from coin.strategy.executor.util.cancel_all_executor import (SimpleCancelAllExecutor)
from coin.strategy.hk_hard_arb.balance_printer3 import CrossTypeArbProducts
from coin.strategy.hk_hard_arb.pricer import (UnhedgePriceCalculator, PositionLeanPricer)
from coin.strategy.hk_hard_arb.util import load_clean_json
from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.feed import SubscriptionRequest, get_flow_book_builder_cls
from coin.strategy.mm.live_envs import (create_live_env_instance, get_live_env_cls)
from coin.strategy.executor.passive.factory import construct_mm_executor
from coin.strategy.executor.passive.util import (convert_executor_config_from_dict)
from coin.support.accounting.checker.account_checker import (AccountChecker,
                                                             AccountCheckFailedError)
from coin.support.accounting.checker.config_util import convert_checker_config_from_dict
from coin.proto.coin_order_gateway_pb2 import OrderGatewayConfig
from coin.exchange.base.order_gateway_config import convert_og_config_from_dict
from coin.exchange.base.strategy_logger import StrategyLogProvider
from coin.proto.coin_request_pb2 import StrategyRequestProto
from coin.proto.coin_query_pb2 import (
    CurrencyBalance,
    AccountBalance,
)
from coin.support.proto_log.logic.pnl_balance_calculator \
    import create_pnl_balance_calculator
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.strategy.hk_hard_arb.util import parse_subreq_config

FLAGS = flags.FLAGS

flags.DEFINE_string('og_configs', '', '')

flags.DEFINE_string('config_filename', '', '')

flags.DEFINE_integer('exit_after_min', 30, '')

flags.DEFINE_boolean('dry_run', False, '')


def _gen_arb_products(common, products, ProductType0, ProductType1):
  if 'trade_quote_currency' in common and 'ref_quote_currency' in common:
    ref_quote_currency = common['ref_quote_currency']
    trade_quote_currency = common['trade_quote_currency']
  elif 'quote_currency' in common:
    ref_quote_currency = common['quote_currency']
    trade_quote_currency = ref_quote_currency
  else:
    raise ValueError(common)

  shared_symbols = []
  symbols_0 = []
  symbols_1 = []
  for shared_symbol, prod_config in products.items():
    if 'ref_symbol' in prod_config:
      symbol_0 = prod_config['ref_symbol']
    else:
      symbol_0 = shared_symbol
    if 'trade_symbol' in prod_config:
      symbol_1 = prod_config['trade_symbol']
    else:
      symbol_1 = shared_symbol
    assert shared_symbol not in shared_symbols, shared_symbol
    assert symbol_0 not in symbols_0
    assert symbol_1 not in symbols_1
    shared_symbols.append(shared_symbol)
    symbols_0.append(symbol_0)
    symbols_1.append(symbol_1)

  arb_products = CrossTypeArbProducts(ref_quote_currency,
                                      trade_quote_currency,
                                      shared_symbols,
                                      symbols_0,
                                      symbols_1,
                                      ProductType0,
                                      ProductType1)
  return arb_products


def _gen_for_same_quotes(config):
  amt_det_notional = config['price_det_notional']
  book_askt_func_0 = (
      lambda book: book.get_notional_asks_by_amt(amt_det_notional).entries[-1].price)
  book_askt_func_1 = book_askt_func_0
  book_bidt_func_0 = (
      lambda book: book.get_notional_bids_by_amt(amt_det_notional).entries[-1].price)
  book_bidt_func_1 = book_bidt_func_0
  return book_askt_func_0, book_askt_func_1, book_bidt_func_0, book_bidt_func_1


def _gen_for_different_quotes(config):
  assert 'price_det_notional' not in config
  assert 'ref_qty_det_notional' in config or 'ref_price_det_notional' in config
  assert 'trade_price_det_notional' in config

  if 'ref_qty_det_notional' in config:
    qty_det_notional = config['ref_qty_det_notional']
    book_askt_func_0 = (
        lambda book: book.get_notional_asks_by_qty(qty_det_notional).entries[-1].price)
    book_bidt_func_0 = (
        lambda book: book.get_notional_bids_by_qty(qty_det_notional).entries[-1].price)
  elif 'ref_price_det_notional' in config:
    amt_det_notional = config['ref_price_det_notional']
    book_askt_func_0 = (
        lambda book: book.get_notional_asks_by_amt(amt_det_notional).entries[-1].price)
    book_bidt_func_0 = (
        lambda book: book.get_notional_bids_by_amt(amt_det_notional).entries[-1].price)
  else:
    raise ValueError('invalie ref')

  amt_det_notional = config['trade_price_det_notional']
  book_askt_func_1 = (
      lambda book: book.get_notional_asks_by_amt(amt_det_notional).entries[-1].price)
  book_bidt_func_1 = (
      lambda book: book.get_notional_bids_by_amt(amt_det_notional).entries[-1].price)
  return book_askt_func_0, book_askt_func_1, book_bidt_func_0, book_bidt_func_1


def _gen_true_book_funcs(config):
  if 'price_det_notional' in config:
    return _gen_for_same_quotes(config)
  else:
    return _gen_for_different_quotes(config)


class SingleProductSubStrategy(object):
  def __init__(self, *, ref_product, trade_product, ccy_0_1, parent, config, logger=None):
    self._trade_product = trade_product
    self._ref_product = ref_product,

    self._logger = logger

    self._edge = None
    self._spread = None
    self._pull_edge = None

    if 'edge_bp' in config:
      self._edge = config['edge_bp'] / 10000.
    if 'spread_bp' in config:
      self._spread = config['spread_bp'] / 10000.
    if 'pull_edge_bp' in config:
      self._pull_edge = config['pull_edge_bp'] / 10000.
    if 'pull_spread_bp' in config:
      self._pull_edge = config['pull_spread_bp'] / 10000. / 2

    assert 1 == sum([
        self._edge is not None,
        self._spread is not None,
    ])

    self._bookmap = BookMap([self._trade_product, self._ref_product], logger=self._logger)

    book_askt_func_0, book_askt_func_1, book_bidt_func_0, book_bidt_func_1 = \
        _gen_true_book_funcs(config)

    self._pricer = UnhedgePriceCalculator(
        basis_ma_window=(int(config['basis_ma_window_min'] * 60e9)),
        tick=get_holder_from_product(trade_product).product_info.gen_tick(),
        book_askt_func_0=book_askt_func_0,
        book_bidt_func_0=book_bidt_func_0,
        book_askt_func_1=book_askt_func_1,
        book_bidt_func_1=book_bidt_func_1,
        ccy_0_1=ccy_0_1)
    self._pos_lean_pricer = PositionLeanPricer(min_pos=config['min_pos'],
                                               max_pos=config['max_pos'],
                                               lot_size=config['lot_size'],
                                               lean_edge_bp=config.get('lean_edge_bp', 0))

    # composite instances
    self._parent = parent
    self._executor = None
    self._config = config

  def on_og_reset(self):
    # Intialize passive executor
    self._logger.info('Intializing passive executor...')
    execution_params = convert_executor_config_from_dict(self._trade_product.symbol, self._config)
    reserve = (execution_params.min_pos + execution_params.max_pos) / 2
    sub_strat_id = self._config.get('sub_strat_id', None)
    self._parent._pnl_balance_calculator.add_product_reserve_sub_strat_id_map(
        self._trade_product, reserve, sub_strat_id)

    self._executor = construct_mm_executor(self._parent._trade_subreq,
                                           execution_params,
                                           self._parent.og,
                                           self._logger)

  def on_book(self, exchange_id, book):
    if exchange_id == 0:
      product = self._ref_product
    elif exchange_id == 1:
      product = self._trade_product
    else:
      raise ValueError(exchange_id)

    self._bookmap.set_book(product, book)

    # Check book freshness and sanity
    if not self._bookmap.is_book_ready() or not self._bookmap.is_book_sane():
      return

    self._pricer.update(self._bookmap.get_book(self._ref_product),
                        self._bookmap.get_book(self._trade_product))

    if not self._pricer.ready or self._executor is None or FLAGS.dry_run:
      return

    if self._spread is not None:
      askp, bidp = self._pricer.calculate_price_by_spread(self._spread)
    elif self._edge is not None:
      askp, bidp = self._pricer.calculate_price(self._edge, self._edge)
    else:
      raise ValueError('not supported mode')

    pos = self._executor.query_unbiased_position()
    askp, bidp = self._pos_lean_pricer.update(pos=pos, buy_price=bidp, sell_price=askp)
    askp, bidp = self._pricer.adjust_price(askp_1=askp, bidp_1=bidp)

    sell_price_pull = None
    buy_price_pull = None
    if askp is not None and self._pull_edge is not None:
      sell_price_pull = askp * (1 - self._pull_edge)
    if bidp is not None and self._pull_edge is not None:
      buy_price_pull = bidp * (1 + self._pull_edge)

    self._executor.manage_mm_orders(sell_price=askp,
                                    buy_price=bidp,
                                    sell_price_pull=sell_price_pull,
                                    buy_price_pull=buy_price_pull)


class PassUnhedgeStrategy(object):
  def __init__(self, strategy_params):
    self._logger = logging.getLogger(__name__)

    self._config = strategy_params['config']

    self._trade_subreq, self._trade_feed_subreq, self._ref_feed_subreq = \
        parse_subreq_config(self._config['common'])

    self._exchange = self._trade_subreq.exchange
    self._og_config = strategy_params['og_config']

    self._balance_checker = None
    account_checker_config = self._get_account_checker_config()
    if account_checker_config is not None:
      self._balance_checker = AccountChecker(account_checker_config)

    # TODO(someone): Need some unified metaclass
    self._book_builder_clses = (
        get_flow_book_builder_cls(self._ref_feed_subreq),
        get_flow_book_builder_cls(self._trade_feed_subreq),
    )
    self._live_env_clses = (
        None,
        get_live_env_cls(self._trade_subreq),
    )
    self._live_env = None

    self.og = None

    self._single_product_sub_strategies = {}
    self._arb_products = _gen_arb_products(self._config['common'],
                                           self._config['products'],
                                           self._book_builder_clses[0].ProductType,
                                           self._book_builder_clses[1].ProductType)

    ref_exchange_feed_checker = self._config['common'].get('ref_exchange_feed_checker', False)
    trade_exchange_feed_checker = self._config['common'].get('trade_exchange_feed_checker', False)
    feed_checker_products = []
    if ref_exchange_feed_checker:
      feed_checker_products += self._arb_products.natives[0].products
    if trade_exchange_feed_checker:
      feed_checker_products += self._arb_products.natives[1].products
    feed_checker_config = get_default_feed_checker_config(feed_checker_products)

    self._feed = MultiFeedSubsystem(feed_checker_config=feed_checker_config)

    self._balance_printer = None

    self._order_canceled_after_og_ready = False
    self._exit_attempt = 0
    self._exit_executor = None

    self._ccy_0_1 = self._config['common'].get('ccy_0_1', None) or 1

    quote_currency = self._config['common']['quote_currency']
    strategy_name = self._config['common']['strategy_name']
    request = StrategyRequestProto(strategy_name=strategy_name)
    self._strategy_log_provider = StrategyLogProvider(request)
    self._pnl_balance_calculator = \
        create_pnl_balance_calculator(self._trade_subreq, quote_currency)

  def _get_account_checker_config(self):
    if self._config.get('account_checker') is None:
      return None

    config_keys = {
        'reserve',
        'lot_size',
        'stack_size',
        'ratio_min_balance',
        'ratio_max_balance',
        'ratio_min_hold_balance',
    }

    balance_check_config_list = []
    for currency, config in self._config['account_checker'].items():
      config = {**config}
      if 'lot_size' not in config:
        for product in self._config['products']:
          base = SharedProduct.FromStrNativeProduct(product).base.currency
          if base == currency:
            config['lot_size'] = self._config['products'][product]['lot_size']
            break

      assert config_keys == set(config.keys()), config
      reserve = config['reserve']
      config['currency'] = currency
      config['min_balance'] = reserve * config['ratio_min_balance']
      config['max_balance'] = reserve * config['ratio_max_balance']
      config['min_available_balance'] = 0.0
      config['max_available_balance'] = config['lot_size'] * config['stack_size']
      config['min_hold_balance'] = reserve * config['ratio_min_hold_balance']
      config['max_hold_balance'] = reserve * config['ratio_max_balance']
      currency_checker_config = convert_checker_config_from_dict(config)
      balance_check_config_list.append(currency_checker_config)

    return AccountCheckerConfig(each_balance_check_config=balance_check_config_list)

  def subscribe_feeds(self):
    common_config = self._config['common']

    for symbol, config in self._config['products'].items():
      config = {**common_config, **config}
      arb_product = self._arb_products.by_shared_symbol(symbol)
      ref_product = arb_product.native_products[0]
      trade_product = arb_product.native_products[1]
      self._single_product_sub_strategies[arb_product.product] = \
        SingleProductSubStrategy(
            ref_product=ref_product,
            trade_product=trade_product,
            ccy_0_1=self._ccy_0_1,
            parent=self,
            config=config,
            logger=self._logger)

    self._feed.subscribe_book(
        FeedSubscriptionRequest.create(self._arb_products.natives[0].products,
                                       self._ref_feed_subreq),
        functools.partial(self._on_book, 0))
    self._feed.subscribe_book(
        FeedSubscriptionRequest.create(self._arb_products.natives[1].products,
                                       self._trade_feed_subreq),
        functools.partial(self._on_book, 1))

  def subscribe_live_envs(self, ioloop, queue, flow_sub):
    og_params = self._config.get('og_params', {})
    og_params = convert_og_config_from_dict(self._trade_subreq, og_params)
    self._live_env = create_live_env_instance(self._live_env_clses[1],
                                              self._arb_products.natives[1].currencies,
                                              self._arb_products.natives[1].products,
                                              self._og_config,
                                              og_params=og_params,
                                              ioloop=ioloop,
                                              queue=queue,
                                              flow_sub=flow_sub)
    self._live_env.on_order_gateway_reset_callback.append(self.register_order_gateway)

  def register_order_gateway(self, name, order_gateway):
    self._logger.info('OG initialized: %s', name)
    self.og = order_gateway
    # cancel working orders first after OGs are ready.
    if not self._order_canceled_after_og_ready:
      self.og.cancel_all()
      self._logger.info('Orders canceled after OGs are ready.')
      self._order_canceled_after_og_ready = True

    for sub_strategy in self._single_product_sub_strategies.values():
      sub_strategy.on_og_reset()

    self._exit_executor = SimpleCancelAllExecutor()

    self._pnl_balance_calculator.add_og(order_gateway)

  def start(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._feed.start(ioloop, queue, flow_sub)
    self._live_env.start()
    self.print_balance(periodic=True)
    PeriodicCallback(self.calc_pnl_balance, 10 * 1000).start()  # every 10 sec

  def _on_book(self, exchange_id, book):
    if not self.og:
      self._logger.info('OG not ready')
      return
    if self._exit_attempt > 0:
      return

    product = self._arb_products.by_exchange_id_and_native_product(exchange_id,
                                                                   book.product).product
    self._single_product_sub_strategies[product].on_book(exchange_id, book)
    self._pnl_balance_calculator.on_book(book)

  def print_balance(self, periodic=False):
    if self._balance_printer is not None:
      self._balance_printer.print_balance()

    if self._balance_checker is not None and self.og is not None and self.og.is_ready():
      balance_proto = self.og.og_info.balance_as_proto()
      try:
        self._balance_checker.check_balances_by_range(balance_proto)
      except AccountCheckFailedError:
        self._logger.exception('Balance is wrong, kill strategy.\n%s', balance_proto)
        os.kill(os.getpid(), signal.SIGINT)  # commit suicide

    if periodic:
      self._ioloop.add_timeout(datetime.timedelta(seconds=10), self.print_balance, True)

  def calc_pnl_balance(self):
    if not self.og:
      return

    pnl_balance_results = self._pnl_balance_calculator.calc_pnl_balance()
    for pnl_balance_result in pnl_balance_results:
      self._strategy_log_provider.write_pnl_balance_result(pnl_balance_result)

  def signal_handler(self, sig, frame):
    self._logger.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self.on_exit)

  def on_exit(self):
    self._exit_attempt += 1
    if self._exit_attempt > 5:
      self._ioloop.add_callback(self._ioloop.stop)

    finished = True
    if self._exit_executor is not None:
      ogs = to_list(self.og)
      self._logger.info('Exit callback executed.')
      finished = self._exit_executor.cancel_all_on_exit(ogs)

    if finished:
      self._ioloop.add_callback(self._ioloop.stop)
    else:
      timeout = datetime.timedelta(seconds=2)
      self._ioloop.add_timeout(timeout, self.on_exit)


def get_param_set(og_configs, config_filename):
  og_config = og_configs.split(',')[0]
  config = load_clean_json(config_filename)
  return {
      'og_config': og_config,
      'config': config,
  }


def main(argv):
  strategy_params = get_param_set(og_configs=FLAGS.og_configs,
                                  config_filename=FLAGS.config_filename)

  queue = SimpleQueue('unhedged_mm')
  flow_sub = flow_subscriber.from_queue(queue)
  ioloop = IOLoop.current()

  strategy = PassUnhedgeStrategy(strategy_params)
  strategy.subscribe_feeds()
  strategy.subscribe_live_envs(ioloop, queue, flow_sub)  # What the..
  strategy.start(ioloop, queue, flow_sub)

  exit_after = datetime.timedelta(minutes=FLAGS.exit_after_min)
  ioloop.add_timeout(exit_after, strategy.on_exit)

  signal.signal(signal.SIGTERM, strategy.signal_handler)
  signal.signal(signal.SIGINT, strategy.signal_handler)

  ioloop.start()


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
